#pragma once
#include <stb_image.h>
#include <stb_image_write.h>
#include <stb_image_resize2.h>
#include <stdexcept>
#include <cassert>
#include "CMacro.h"

class CImage {
public:
	CImage() : m_isSTBI(false), m_width(0), m_height(0), m_channels(0), m_bytes(nullptr) {}

	CImage(const std::string& path) : m_isSTBI(true) {
		m_bytes = load(path, &m_width, &m_height, &m_channels);
	}

	CImage(UChar* compressedRawBytes, int len) : m_isSTBI(true) {
		m_bytes = load(compressedRawBytes, len, &m_width, &m_height, &m_channels);
	}

	CImage(CByteArray& compressedRawBytes) : m_isSTBI(true) {
		m_bytes = load(compressedRawBytes.data(), compressedRawBytes.size(), &m_width, &m_height, &m_channels);
	}

	CImage(UChar* uncompressedBytes, int width, int height, int channels) : m_isSTBI(false), m_bytes(uncompressedBytes), m_width(width), m_height(height), m_channels(channels) {}

	CImage(const CByteArray& uncompressedByteArray, int width, int height, int channels) : m_isSTBI(false), m_width(width), m_height(height), m_channels(channels) {
		size_t size = m_width * m_height * m_channels;
		m_bytes = new UChar[size];
		std::memcpy(m_bytes, uncompressedByteArray.data(), size);
	}

	CImage(const CImage& other) : m_isSTBI(other.m_isSTBI), m_width(other.m_width), m_height(other.m_height), m_channels(other.m_channels) {
		m_bytes = new UChar[m_width * m_height * m_channels];
		std::memcpy(m_bytes, other.m_bytes, m_width * m_height * m_channels);
	}

	CImage(CImage&& other) noexcept : m_isSTBI(other.m_isSTBI), m_width(other.m_width), m_height(other.m_height), m_channels(other.m_channels), m_bytes(other.m_bytes) {
		other.m_bytes = nullptr;
		other.m_width = other.m_height = other.m_channels = 0;
	}

	~CImage() {
		if (m_isSTBI)
			stbi_image_free(m_bytes);
		else
			delete[] m_bytes;
	}

	CImage& operator=(const CImage& other) {
		if (this != &other) {
			delete[] m_bytes;
			m_isSTBI = other.m_isSTBI;
			m_width = other.m_width;
			m_height = other.m_height;
			m_channels = other.m_channels;
			m_bytes = new UChar[m_width * m_height * m_channels];
			std::memcpy(m_bytes, other.m_bytes, m_width * m_height * m_channels);
		}
		return *this;
	}

	CImage& operator=(CImage&& other) noexcept {
		if (this != &other) {
			delete[] m_bytes;
			m_isSTBI = other.m_isSTBI;
			m_width = other.m_width;
			m_height = other.m_height;
			m_channels = other.m_channels;
			m_bytes = other.m_bytes;
			other.m_bytes = nullptr;
		}
		return *this;
	}

	/** also good
	CImage& resize(int new_width, int new_height) {// bilinear interpolation
		assert(m_bytes != nullptr);
		if (m_width == new_width && m_height == new_height)
			return *this;

		CByteArray new_data(new_width * new_height * m_channels);

		for (int y = 0; y < new_height; ++y) {
			for (int x = 0; x < new_width; ++x) {
				// Calculate the corresponding position in the source image
				float src_x = (x + 0.5f) * m_width / new_width - 0.5f;
				float src_y = (y + 0.5f) * m_height / new_height - 0.5f;
				// Calculate the coordinates of four adjacent pixels
				int x1 = static_cast<int>(src_x);
				int y1 = static_cast<int>(src_y);
				int x2 = std::min(x1 + 1, m_width - 1);
				int y2 = std::min(y1 + 1, m_height - 1);
				// Calculate interpolation weights
				float wx = src_x - x1;
				float wy = src_y - y1;
				for (int c = 0; c < m_channels; ++c) {
					// Bilinear interpolation
					float val = (1 - wx) * (1 - wy) * m_bytes[(y1 * m_width + x1) * m_channels + c] +
								wx * (1 - wy) * m_bytes[(y1 * m_width + x2) * m_channels + c] +
								(1 - wx) * wy * m_bytes[(y2 * m_width + x1) * m_channels + c] +
								wx * wy * m_bytes[(y2 * m_width + x2) * m_channels + c];

					new_data[(y * new_width + x) * m_channels + c] = static_cast<UChar>(std::round(val));
				}
			}
		}

		delete[] m_bytes;
		m_bytes = new UChar[new_width * new_height * m_channels];
		std::memcpy(m_bytes, new_data.data(), new_width * new_height * m_channels);
		m_width = new_width;
		m_height = new_height;
		return *this;
	}*/

	CImage& resize(int newWidth, int newHeight) {
		assert(m_bytes != nullptr);
		if (m_width == newWidth && m_height == newHeight)
			return *this;

		auto resizedBytes = new UChar[newWidth * newHeight * m_channels];
		stbir_pixel_layout pixelLayout = m_channels >= 4 ? STBIR_RGBA : STBIR_RGB;
		stbir_resize_uint8_linear(m_bytes, m_width, m_height, 0, resizedBytes, newWidth, newHeight, 0, pixelLayout);
		
		stbi_image_free(m_bytes);
		m_bytes = resizedBytes;
		m_width = newWidth;
		m_height = newHeight;

		return *this;
	}

	CImage& invertY() {// Flip the image vertically
		assert(m_bytes != nullptr);
		for (int y = 0; y < m_height / 2; ++y) {
			for (int x = 0; x < m_width; ++x) {
				for (int c = 0; c < m_channels; ++c) {
					std::swap(m_bytes[(y * m_width + x) * m_channels + c],
							  m_bytes[((m_height - 1 - y) * m_width + x) * m_channels + c]);
				}
			}
		}
		return *this;
	}

	CImage& crop(int x, int y, int width, int height) {
		assert(m_bytes != nullptr);
		if (x < 0 || y < 0 || width <= 0 || height <= 0 || x + width > m_width || y + height > m_height)
			throw std::invalid_argument("Invalid crop region");
		// allocate memory for the cropped image
		size_t newSize = width * height * m_channels;
		UChar* newBytes = new UChar[newSize];
		// copy the cropped region from the original image
		for (int row = 0; row < height; ++row) {
			for (int col = 0; col < width; ++col) {
				// calculate the source and destination indices
				size_t srcIndex = ((y + row) * m_width + (x + col)) * m_channels;
				size_t destIndex = (row * width + col) * m_channels;
				// copy the pixel data
				for (int c = 0; c < m_channels; ++c)
					newBytes[destIndex + c] = m_bytes[srcIndex + c];
			}
		}
		// release the old memory and update the image data
		delete[] m_bytes;
		m_bytes = newBytes;
		// update the image dimensions
		m_width = width;
		m_height = height;
		return *this;
	}

	bool save(const std::string& path, int jpegQuality = 100) {
		if (m_bytes == nullptr || m_width <= 0 || m_height <= 0 || m_channels <= 0)
			return false;
		int result = 0;
		if (m_channels == 3)
			result = stbi_write_jpg(path.c_str(), m_width, m_height, m_channels, m_bytes, jpegQuality);
		else if (m_channels == 4)
			result = stbi_write_png(path.c_str(), m_width, m_height, m_channels, m_bytes, m_width * m_channels);
		else
			return false;// unsupported channels
		// use `stbi_write_bmp` to save uncompressed image
		return result != 0;
	}

	int width() const { return m_width; }
	int height() const { return m_height; }
	int channels() const { return m_channels; }
	const UChar* data() const { return m_bytes; }

	bool isValid() const { return m_bytes != nullptr && m_width > 0 && m_height > 0; }

private:
	/**
	 *@param rawBytes must have header data, note that it is compressed(JPEG/PNG), some BMP format is raw uncompressed!
	 *@warning release it manually!
	 *@return uncompressed bytes
	 */
	static UChar* load(UChar* compressedRawBytes, int len, int* width, int* height, int* channels) {
		return stbi_load_from_memory(compressedRawBytes, len, width, height, channels, 0);
	}

	/**
	 *@warning release it manually!
	 *@return uncompressed bytes
	 */
	static UChar* load(const std::string& path, int* width, int* height, int* channels) {
		return stbi_load(path.c_str(), width, height, channels, 0);
	}

	static bool info(UChar* bytes, int len, int* width, int* height, int* channels) {
		return stbi_info_from_memory(bytes, len, width, height, channels) == 1;
	}

	static bool info(const std::string& path, int* width, int* height, int* channels) {
		return stbi_info(path.c_str(), width, height, channels) == 1;
	}

	static void releaseStbiImage(UChar* bytes) { stbi_image_free(bytes); }

	static void printFailureReasonMsg() { std::cerr << "stbi failure reason is: " << stbi_failure_reason() << std::endl; }


	/**@important uncompressed bytes!*/
	UChar* m_bytes;

	int m_width, m_height;
	int m_channels;

	bool m_isSTBI;
};